# -*- coding: utf-8 -*-
"""
均线排列股票筛选器
筛选规则：
1. 最近10个交易日，5日均线都在10日均线之上
2. 最后一个交易日：5日均线 > 10日均线 > 20日均线 > 30日均线
3. 最近10个交易日的最大涨幅不超过5%

这种均线排列的情况可能预示着：
1. 股票处于稳健上升趋势
2. 多头排列，趋势明确
3. 涨幅温和，没有过度炒作
4. 潜在的趋势延续机会

使用方法：
    python ma_alignment_screener.py              # 筛选所有股票
    python ma_alignment_screener.py --limit 100  # 限制筛选数量（测试用）
    python ma_alignment_screener.py --stock 000001  # 筛选指定股票
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
import argparse

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'ma_alignment_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class MAAlignmentScreener:
    """均线排列股票筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)
    
    def get_stock_codes(self, limit: int = None) -> list:
        """获取所有活跃股票代码"""
        try:
            sql = "SELECT stock_code, stock_name FROM new_stock_info WHERE status = 'active' ORDER BY stock_code"
            if limit:
                sql += f" LIMIT {limit}"
            
            results = self.db.execute_query(sql)
            return [(row['stock_code'], row['stock_name']) for row in results]
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_daily_data(self, stock_code: str, days: int = 50) -> pd.DataFrame:
        """
        获取股票最近N天的日线数据
        
        Args:
            stock_code: 股票代码
            days: 获取天数（默认50天，确保有足够的数据计算30日均线）
            
        Returns:
            包含交易数据的DataFrame
        """
        sql = """
        SELECT trade_date, close_price, change_pct
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换数值列
            df['close_price'] = pd.to_numeric(df['close_price'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            
            # 按日期升序排列
            df = df.sort_values('trade_date').reset_index(drop=True)
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_moving_averages(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算移动平均线
        
        Args:
            df: 股票日线数据DataFrame
            
        Returns:
            添加了均线列的DataFrame
        """
        if df.empty:
            return df
        
        # 计算各周期均线
        df['ma5'] = df['close_price'].rolling(window=5).mean()
        df['ma10'] = df['close_price'].rolling(window=10).mean()
        df['ma20'] = df['close_price'].rolling(window=20).mean()
        df['ma30'] = df['close_price'].rolling(window=30).mean()
        
        return df
    
    def check_ma_alignment(self, df: pd.DataFrame) -> dict:
        """
        检查是否满足均线排列条件
        
        条件：
        1. 最近10个交易日，5日均线都在10日均线之上
        2. 最后一个交易日：5日均线 > 10日均线 > 20日均线 > 30日均线
        3. 最近10个交易日的最大涨幅不超过5%
        
        Args:
            df: 股票日线数据DataFrame（已计算均线）
            
        Returns:
            符合条件返回详细信息字典，否则返回None
        """
        if df.empty or len(df) < 30:  # 至少需要30天数据来计算30日均线
            return None
        
        # 检查是否有空值
        if df[['ma5', 'ma10', 'ma20', 'ma30']].isna().any().any():
            return None
        
        # 获取最近10个交易日的数据
        recent_10d = df.tail(10).copy()
        
        if len(recent_10d) < 10:
            return None
        
        # 条件1: 最近10个交易日，5日均线都在10日均线之上
        ma5_above_ma10 = (recent_10d['ma5'] > recent_10d['ma10']).all()
        
        if not ma5_above_ma10:
            return None
        
        # 获取最后一个交易日的数据
        latest_row = df.iloc[-1]
        
        # 条件2: 最后一个交易日：5日均线 > 10日均线 > 20日均线 > 30日均线
        ma_perfect_alignment = (
            latest_row['ma5'] > latest_row['ma10'] > 
            latest_row['ma20'] > latest_row['ma30']
        )
        
        if not ma_perfect_alignment:
            return None
        
        # 条件3: 最近10个交易日的最大涨幅不超过5%
        max_gain = recent_10d['change_pct'].max()
        
        if max_gain > 5.0:
            return None
        
        # 计算额外的统计信息
        # 计算均线之间的距离（作为趋势强度指标）
        ma5_ma10_gap = ((latest_row['ma5'] - latest_row['ma10']) / latest_row['ma10'] * 100)
        ma10_ma20_gap = ((latest_row['ma10'] - latest_row['ma20']) / latest_row['ma20'] * 100)
        ma20_ma30_gap = ((latest_row['ma20'] - latest_row['ma30']) / latest_row['ma30'] * 100)
        
        # 计算最近10天的平均涨幅和总涨幅
        avg_change = recent_10d['change_pct'].mean()
        total_change = ((recent_10d.iloc[-1]['close_price'] - recent_10d.iloc[0]['close_price']) 
                       / recent_10d.iloc[0]['close_price'] * 100)
        
        # 计算最近10天的最小涨幅
        min_change = recent_10d['change_pct'].min()
        
        return {
            'trade_date': latest_row['trade_date'],
            'close_price': latest_row['close_price'],
            'ma5': round(latest_row['ma5'], 2),
            'ma10': round(latest_row['ma10'], 2),
            'ma20': round(latest_row['ma20'], 2),
            'ma30': round(latest_row['ma30'], 2),
            'ma5_ma10_gap': round(ma5_ma10_gap, 2),
            'ma10_ma20_gap': round(ma10_ma20_gap, 2),
            'ma20_ma30_gap': round(ma20_ma30_gap, 2),
            'max_gain_10d': round(max_gain, 2),
            'min_change_10d': round(min_change, 2),
            'avg_change_10d': round(avg_change, 2),
            'total_change_10d': round(total_change, 2),
            'latest_change_pct': round(latest_row['change_pct'], 2) if pd.notna(latest_row['change_pct']) else 0
        }
    
    def screen_stocks(self, stock_codes: list = None, limit: int = None) -> pd.DataFrame:
        """
        筛选符合条件的股票
        
        Args:
            stock_codes: 指定要筛选的股票列表，格式为[(code, name), ...]
            limit: 限制筛选的股票数量
            
        Returns:
            包含筛选结果的DataFrame
        """
        if stock_codes is None:
            stock_codes = self.get_stock_codes(limit=limit)
        
        if not stock_codes:
            logger.warning("没有找到需要筛选的股票")
            return pd.DataFrame()
        
        logger.info(f"开始筛选 {len(stock_codes)} 只股票...")
        
        results = []
        total = len(stock_codes)
        
        for idx, (stock_code, stock_name) in enumerate(stock_codes, 1):
            try:
                # 每10只股票输出一次进度
                if idx % 10 == 0 or idx == total:
                    logger.info(f"进度: {idx}/{total} ({idx/total*100:.1f}%)")
                
                # 获取日线数据
                df = self.get_stock_daily_data(stock_code, days=50)
                
                if df.empty or len(df) < 30:
                    continue
                
                # 计算均线
                df = self.calculate_moving_averages(df)
                
                # 检查是否满足条件
                result = self.check_ma_alignment(df)
                
                if result:
                    result['stock_code'] = stock_code
                    result['stock_name'] = stock_name
                    results.append(result)
                    
                    logger.info(
                        f"✓ {stock_code} {stock_name} - "
                        f"价格: {result['close_price']:.2f}, "
                        f"MA(5/10/20/30): {result['ma5']:.2f}/{result['ma10']:.2f}/"
                        f"{result['ma20']:.2f}/{result['ma30']:.2f}, "
                        f"10日最大涨幅: {result['max_gain_10d']:.2f}%"
                    )
            
            except Exception as e:
                logger.error(f"处理股票 {stock_code} {stock_name} 时出错: {e}")
                continue
        
        if not results:
            logger.info("未找到符合条件的股票")
            return pd.DataFrame()
        
        # 转换为DataFrame并排序
        df_results = pd.DataFrame(results)
        
        # 按总涨幅降序排列（涨幅越大越靠前）
        df_results = df_results.sort_values('total_change_10d', ascending=False)
        
        # 重新排列列的顺序
        columns_order = [
            'stock_code', 'stock_name', 'trade_date', 'close_price',
            'ma5', 'ma10', 'ma20', 'ma30',
            'ma5_ma10_gap', 'ma10_ma20_gap', 'ma20_ma30_gap',
            'latest_change_pct', 'total_change_10d', 'avg_change_10d',
            'max_gain_10d', 'min_change_10d'
        ]
        df_results = df_results[columns_order]
        
        logger.info(f"\n筛选完成！共找到 {len(df_results)} 只符合条件的股票")
        
        return df_results
    
    def save_results(self, df: pd.DataFrame, output_dir: str = '.') -> tuple:
        """
        保存筛选结果到CSV和TXT文件
        
        Args:
            df: 筛选结果DataFrame
            output_dir: 输出目录
            
        Returns:
            (csv_file, txt_file) 文件路径元组
        """
        if df.empty:
            logger.warning("没有结果需要保存")
            return None, None
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        csv_file = os.path.join(output_dir, f'ma_alignment_stocks_{timestamp}.csv')
        txt_file = os.path.join(output_dir, f'stock_codes_ma_alignment_{timestamp}.txt')
        
        try:
            # 保存CSV文件
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"结果已保存到: {csv_file}")
            
            # 保存股票代码到TXT文件
            with open(txt_file, 'w', encoding='utf-8') as f:
                for _, row in df.iterrows():
                    f.write(f"{row['stock_code']}\n")
            logger.info(f"股票代码已保存到: {txt_file}")
            
            return csv_file, txt_file
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
            return None, None
    
    def print_summary(self, df: pd.DataFrame):
        """打印筛选结果摘要"""
        if df.empty:
            return
        
        logger.info("\n" + "="*80)
        logger.info("筛选结果摘要")
        logger.info("="*80)
        logger.info(f"符合条件的股票数量: {len(df)}")
        logger.info(f"平均10日涨幅: {df['total_change_10d'].mean():.2f}%")
        logger.info(f"最大10日涨幅: {df['total_change_10d'].max():.2f}%")
        logger.info(f"最小10日涨幅: {df['total_change_10d'].min():.2f}%")
        logger.info(f"平均MA5-MA10间距: {df['ma5_ma10_gap'].mean():.2f}%")
        
        logger.info("\n前10只涨幅最大的股票:")
        logger.info("-"*80)
        for idx, row in df.head(10).iterrows():
            logger.info(
                f"{row['stock_code']} {row['stock_name']:8s} - "
                f"价格: {row['close_price']:7.2f} | "
                f"10日涨幅: {row['total_change_10d']:6.2f}% | "
                f"10日最大涨: {row['max_gain_10d']:5.2f}% | "
                f"MA5: {row['ma5']:7.2f}"
            )
        logger.info("="*80)
        
        # 输出逗号分隔的股票代码
        stock_codes = df['stock_code'].tolist()
        stock_codes_str = ','.join(stock_codes)
        
        logger.info("\n" + "="*80)
        logger.info("所有匹配的股票代码（逗号分隔）:")
        logger.info("="*80)
        logger.info(stock_codes_str)
        logger.info("="*80 + "\n")


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='均线排列股票筛选器 - 筛选均线排列良好且涨幅温和的股票'
    )
    parser.add_argument(
        '--limit',
        type=int,
        help='限制筛选的股票数量（用于测试）'
    )
    parser.add_argument(
        '--stock',
        type=str,
        help='筛选指定股票代码'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='输出文件目录（默认：当前目录）'
    )
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_arguments()
    
    logger.info("="*80)
    logger.info("均线排列股票筛选器")
    logger.info("筛选规则:")
    logger.info("1. 最近10个交易日，5日均线都在10日均线之上")
    logger.info("2. 最后一个交易日：5日均线 > 10日均线 > 20日均线 > 30日均线")
    logger.info("3. 最近10个交易日的最大涨幅不超过5%")
    logger.info("="*80 + "\n")
    
    screener = MAAlignmentScreener()
    
    # 如果指定了股票代码，只筛选该股票
    if args.stock:
        stock_info = screener.stock_dao.get_stock_info_by_code(args.stock)
        if not stock_info:
            logger.error(f"未找到股票: {args.stock}")
            return
        
        stock_codes = [(stock_info['stock_code'], stock_info['stock_name'])]
        logger.info(f"筛选指定股票: {args.stock} {stock_info['stock_name']}\n")
    else:
        stock_codes = None
    
    # 执行筛选
    start_time = datetime.now()
    df_results = screener.screen_stocks(stock_codes=stock_codes, limit=args.limit)
    end_time = datetime.now()
    
    # 打印结果摘要
    screener.print_summary(df_results)
    
    # 保存结果
    if not df_results.empty:
        csv_file, txt_file = screener.save_results(df_results, output_dir=args.output_dir)
        
        if csv_file and txt_file:
            logger.info(f"\n筛选用时: {(end_time - start_time).total_seconds():.2f} 秒")
            logger.info(f"详细结果请查看: {csv_file}")
            logger.info(f"股票代码列表: {txt_file}")
    
    logger.info("\n筛选完成！")


if __name__ == '__main__':
    main()
